package org.ayachinene.leetcode.problems;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class LCP07 {

    static class Solution {

        /*
         * key: 目标小朋友的编号
         * value: 可以把信息传给目标小朋友的小朋友们的编号
         */
        private final Map<Integer, List<Integer>> source = new HashMap<>();

        public int numWays(int n, int[][] relation, int k) {
            // 初始化source
            for (int[] r : relation) {
                int src = r[0];
                int tar = r[1];

                if (!source.containsKey(tar)) {
                    source.put(tar, new ArrayList<>());
                }
                source.get(tar).add(src);
            }

            return numWays(n - 1, k);
        }

        // 经k轮将信息从[编号为0的小朋友]传递给[编号为id的小朋友]的方案数
        private int numWays(int id, int k) {
            if (!source.containsKey(id)) {
                return 0;
            }

            if (k == 1) {
                return (int) source.get(id).stream()
                        .filter(i -> i == 0)
                        .count();
            }

            int res = 0;
            for (Integer srcId : source.get(id)) {
                res += numWays(srcId, k - 1);
            }
            return res;
        }
    }
}
